from PyQt5.QtCore import Qt, QRect
from PyQt5.QtWidgets import QApplication, QPushButton, QLabel, QMainWindow, QFileDialog, QLineEdit, QAction, \
	QDesktopWidget
from PyQt5.QtGui import QPixmap, QPen, QPainter
import sys
import os
from xml.dom import minidom


class VOC_Sample_Generator:

    def __init__(self):
        self.dom = minidom.Document()
        self.root_node = self.dom.createElement('annotation')
        self.object_node = self.dom.createElement('object')
        self.root_node.appendChild(self.object_node)

    def add_filename(self, filename):
        text = self.dom.createTextNode(filename)
        filename_node = self.dom.createElement('filename')
        filename_node.appendChild(text)
        self.root_node.appendChild(filename_node)

    def add_class(self, class_name):
        text = self.dom.createTextNode(class_name)
        name_node = self.dom.createElement('name')
        name_node.appendChild(text)
        self.object_node.appendChild(name_node)

    def add_size(self, width, height, depth):
        text = self.dom.createTextNode(str(width))
        width_node = self.dom.createElement('width')
        width_node.appendChild(text)

        text = self.dom.createTextNode(str(height))
        height_node = self.dom.createElement('height')
        height_node.appendChild(text)

        text = self.dom.createTextNode(str(depth))
        depth_node = self.dom.createElement('depth')
        depth_node.appendChild(text)

        size_node = self.dom.createElement('size')

        size_node.appendChild(width_node)
        size_node.appendChild(height_node)
        size_node.appendChild(depth_node)
        self.root_node.appendChild(size_node)

    def add_bndbox(self, xmin, ymin, xmax, ymax):
        text = self.dom.createTextNode(str(xmin))
        xmin_node = self.dom.createElement('xmin')
        xmin_node.appendChild(text)

        text = self.dom.createTextNode(str(ymin))
        ymin_node = self.dom.createElement('ymin')
        ymin_node.appendChild(text)

        text = self.dom.createTextNode(str(xmax))
        xmax_node = self.dom.createElement('xmax')
        xmax_node.appendChild(text)

        text = self.dom.createTextNode(str(ymax))
        ymax_node = self.dom.createElement('ymax')
        ymax_node.appendChild(text)

        bndbox_node = self.dom.createElement('bndbox')

        bndbox_node.appendChild(xmin_node)
        bndbox_node.appendChild(ymin_node)
        bndbox_node.appendChild(xmax_node)
        bndbox_node.appendChild(ymax_node)
        self.object_node.appendChild(bndbox_node)

    def build(self, path):
        self.dom.appendChild(self.root_node)
        with open(path, 'w', encoding='utf-8') as f:
            self.dom.writexml(f, indent='', addindent='\t', newl='\n', encoding='utf-8')
        f.close()


class OtherLabel(QLabel):

	def __init__(self, parent=None):
		super(OtherLabel, self).__init__(parent)
		self.x0 = 0
		self.x1 = 0
		self.y0 = 0
		self.y1 = 0
		self.flag = False
		self.move = False

	def mouseMoveEvent(self, event):
		self.move = True
		if self.flag:
			self.x1 = event.pos().x()
			self.y1 = event.pos().y()
			self.update()

	def mouseReleaseEvent(self, event):
		self.flag = False
		self.move = False

	def mousePressEvent(self, event):
		self.x0 = event.pos().x()
		self.y0 = event.pos().y()
		self.flag = True


class MainLabel(QLabel):
	def __init__(self, parent=None):
		super(MainLabel, self).__init__(parent)
		self.x0 = 0
		self.y0 = 0
		self.x1 = 0
		self.y1 = 0
		self.flag = False
		self.move = False

	def mouseMoveEvent(self, event):
		self.move = True
		if self.flag:
			self.x1 = event.pos().x()
			self.y1 = event.pos().y()
			self.update()

	def mouseReleaseEvent(self, event):
		self.flag = False
		self.move = False

	def paintEvent(self, event):
		super().paintEvent(event)
		if self.flag and self.move:
			rect = QRect(self.x0, self.y0, abs(self.x1 - self.x0), abs(self.y1 - self.y0))
			painter = QPainter(self)
			painter.setPen(QPen(Qt.red, 2, Qt.SolidLine))
			painter.drawRect(rect)

	def mousePressEvent(self, event):
		self.x0 = event.pos().x()
		self.y0 = event.pos().y()
		self.flag = True


class mainForm(QMainWindow):
	def __init__(self):
		super(mainForm, self).__init__()

		toolBar = self.addToolBar("File")
		new = QAction("文件", self)
		toolBar.addAction(new)
		new.triggered.connect(self.fileopen)
		toolBar.setToolButtonStyle(Qt.ToolButtonTextOnly)

		self.curshow = QLabel("", self)
		self.curshow.move(550, 400)
		self.curshow.setFixedHeight(30)
		self.curshow.setFixedWidth(100)

		self.setWindowTitle("数据集标准工具")
		self.resize(900, 500)

	def center(self):
		screen = QDesktopWidget().screenGeometry()
		size = self.geometry()
		newLeft = (screen.width() - size.width()) / 2
		newTop = (screen.height() - size.height()) / 2
		self.move(newLeft, newTop)

	def fileopen(self):
		self.img_path = QFileDialog.getExistingDirectory(self, '选取文件夹', './')
		if self.img_path == '':
			return
		self.img_list = os.listdir(self.img_path)  # 获取目录下的所有文件
		tmp = []
		for i in self.img_list:
			if i.split('.')[-1] == 'jpg' or i.split('.')[-1] == 'png':
				tmp.append(i)
		self.img_list = tmp
		self.idx = 0
		self.initUI()


	def initUI(self):

		self.bt = QPushButton('确认', self)
		self.bt.move(130, 400)
		self.bt.setFixedHeight(30)
		self.bt.clicked.connect(self.classify)

		self.text = QLabel('标注', self)
		self.text.move(270, 400)
		self.text.setFixedHeight(30)

		self.editline = QLineEdit(self)
		self.editline.move(300, 400)
		self.editline.setFixedHeight(30)
		self.editline.setFixedWidth(200)

		self.tooltip = QLabel("", self)
		self.tooltip.move(300, 450)
		self.tooltip.setFixedHeight(30)

		self.lbl_list = []

		for i in range(self.get_remainder()):  #
			self.pix = QPixmap(self.img_path + '/' + self.img_list[self.idx + i])
			if i == 0:
				self.label_img = MainLabel(self)
			else:
				self.label_img = OtherLabel(self)
			if i == 0:
				self.label_img.setGeometry(360 * i + 10, 300 - 100 * (3 - i) + 40, 100 * (3 - i) + 40, 100 * (3 - i) + 40)
			elif i == 1:
				self.label_img.setGeometry(360 * i + 100, 300 - 100 + 40, 100 + 40, 100 + 40)
			else:
				self.label_img.setGeometry(360 * i + 10, 300 - 100 + 40, 100 + 40, 100 + 40)
			self.label_img.setStyleSheet("border:2px solid black")
			self.label_img.setPixmap(self.pix)
			self.label_img.setScaledContents(True)
			self.lbl_list.append(self.label_img)
			self.label_img.show()
		self.bt.show()
		self.editline.show()
		self.tooltip.show()
		self.text.show()
		self.curshow.show()
		self.curshow.setText('')
		if len(self.lbl_list) != 0:
			self.curshow.setText("当前是第%d张图" % (self.idx + 1))
		else:
			self.curshow.setText('此文件夹无图')

	def get_remainder(self):
		r = len(self.img_list) - self.idx
		if r > 3:
			r = 3
		return r

	def clear_lbls(self):
		for i in range(len(self.lbl_list)):
			self.lbl_list[i].hide()

	def classify(self):
		if self.editline.text() == '':
			self.tooltip.setText('标注不能为空！')
			return
		elif self.lbl_list[0].x0 == self.lbl_list[0].x1 and self.lbl_list[0].y0 == self.lbl_list[0].y1:
			self.tooltip.setText('未进行框选！')
			return
		else:
			self.tooltip.setText('')

		voc = VOC_Sample_Generator()
		filename = self.img_list[self.idx].split('.')[0]
		voc.add_bndbox(self.lbl_list[0].x0, self.lbl_list[0].y0, self.lbl_list[0].x1, self.lbl_list[0].y1)
		voc.add_filename(filename + '.xml')
		voc.add_size(self.label_img.width(), self.label_img.height(), 3)
		voc.add_class(self.editline.text())
		self.editline.setText('')
		voc.build(self.img_path + '/' + 'outputs/' + filename + '.xml')

		self.idx += 1
		img_full_path = [self.img_path + '/' + self.img_list[self.idx + i] for i in range(self.get_remainder())]

		self.clear_lbls()
		for i in range(self.get_remainder()):
			pix = QPixmap(img_full_path[i])
			self.lbl_list[i].setPixmap(pix)
			self.lbl_list[i].show()

		if len(self.lbl_list) != 0:
			self.curshow.setText("当前是第%d张图" % (self.idx + 1))
		else:
			self.curshow.setText('')


app = QApplication(sys.argv)
f = mainForm()
f.show()
sys.exit(app.exec())
